import psutil
import time
import argparse
import os

def get_system_info():
    """获取系统信息"""
    # CPU信息
    cpu_percent = psutil.cpu_percent(interval=1)
    cpu_count = psutil.cpu_count(logical=False)
    cpu_count_logical = psutil.cpu_count(logical=True)
    
    # 内存信息
    mem = psutil.virtual_memory()
    mem_total = mem.total
    mem_used = mem.used
    mem_percent = mem.percent
    
    # 磁盘信息（根目录）
    disk = psutil.disk_usage('/')
    disk_total = disk.total
    disk_used = disk.used
    disk_percent = disk.percent
    
    # 网络信息
    net_io = psutil.net_io_counters()
    net_sent = net_io.bytes_sent
    net_recv = net_io.bytes_recv
    
    return {
        'cpu': {
            'percent': cpu_percent,
            'physical_cores': cpu_count,
            'logical_cores': cpu_count_logical
        },
        'memory': {
            'total': mem_total,
            'used': mem_used,
            'percent': mem_percent
        },
        'disk': {
            'total': disk_total,
            'used': disk_used,
            'percent': disk_percent
        },
        'network': {
            'sent': net_sent,
            'recv': net_recv
        }
    }

def format_bytes(size):
    """格式化字节为易读单位"""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size < 1024.0:
            return f"{size:.2f} {unit}"
        size /= 1024.0
    return f"{size:.2f} PB"

def monitor_system(interval=5, duration=None):
    """监控系统信息"""
    start_time = time.time()
    try:
        print("系统资源监控 (按Ctrl+C停止)\n")
        while True:
            # 清空屏幕
            os.system('cls' if os.name == 'nt' else 'clear')
            
            # 获取并显示系统信息
            info = get_system_info()
            
            print(f"监控时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
            print("-" * 50)
            
            # CPU信息
            print(f"CPU使用率: {info['cpu']['percent']}%")
            print(f"CPU核心数: 物理={info['cpu']['physical_cores']}, 逻辑={info['cpu']['logical_cores']}")
            print("-" * 50)
            
            # 内存信息
            print(f"内存使用率: {info['memory']['percent']}%")
            print(f"内存使用: {format_bytes(info['memory']['used'])} / {format_bytes(info['memory']['total'])}")
            print("-" * 50)
            
            # 磁盘信息
            print(f"磁盘使用率: {info['disk']['percent']}%")
            print(f"磁盘使用: {format_bytes(info['disk']['used'])} / {format_bytes(info['disk']['total'])}")
            print("-" * 50)
            
            # 网络信息
            print(f"网络发送: {format_bytes(info['network']['sent'])}")
            print(f"网络接收: {format_bytes(info['network']['recv'])}")
            print("-" * 50)
            
            # 检查是否达到监控时长
            if duration and (time.time() - start_time) >= duration:
                print(f"\n监控结束（已运行 {duration} 秒）")
                break
                
            time.sleep(interval)
            
    except KeyboardInterrupt:
        print("\n监控已停止")
    except Exception as e:
        print(f"监控出错: {str(e)}")

if __name__ == "__main__":
    # 安装依赖提示
    try:
        import psutil
    except ImportError:
        print("请先安装依赖: pip install psutil")
        exit(1)
    
    parser = argparse.ArgumentParser(description='系统资源监控脚本')
    parser.add_argument('--interval', type=int, default=5, help='刷新间隔（秒），默认5秒')
    parser.add_argument('--duration', type=int, help='监控时长（秒），不指定则持续监控')
    
    args = parser.parse_args()
    monitor_system(args.interval, args.duration)
